package cn.xaut.ocr3.service.impl;

import cn.xaut.ocr3.dto.IdCard;
import cn.xaut.ocr3.service.RapidOcrService;
import cn.xaut.ocr3.utils.RapidOcrUtils;
import io.github.mymonstercat.ocr.InferenceEngine;
import io.github.mymonstercat.ocr.config.ParamConfig;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.PostConstruct;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author dgl
 */
@Service
public class RapidOcrServiceImpl implements RapidOcrService {

    private final InferenceEngine engine;
    private final ParamConfig paramConfig;

    public RapidOcrServiceImpl(InferenceEngine engine, ParamConfig paramConfig) {
        this.engine = engine;
        this.paramConfig = paramConfig;
    }

    String originPath;
    @PostConstruct
    public void initOriginPath(){
        //获取临时文件路径
        originPath = System.getProperty("java.io.tmpdir") + File.separator + "ocrFile";
    }

    /**
     * 识别图片
     * @param file
     * @return
     * @throws IOException
     */
    @Override
    public String recognizeImage(MultipartFile file) throws IOException {

        List<String> list = new ArrayList<>();
        String fileName = file.getOriginalFilename();

        List<String> imageNames = Arrays.asList("jpg", "png", "jpeg");
        List<String> pdfNames = Arrays.asList("pdf","PDF");
        String endFileName = fileName.substring(fileName.lastIndexOf(".") + 1).toLowerCase();

        //判断文件格式
        if (pdfNames.contains(endFileName)) {

            new File(originPath).mkdirs();
            File tempFile = new File(originPath + File.separator + file.getOriginalFilename());
            file.transferTo(tempFile);

            //将pdf文件转为图片
            PDDocument document = PDDocument.load(tempFile);
            PDFRenderer pdf =  new PDFRenderer(document);
            int pageCount = document.getNumberOfPages();

            //识别每一页
            for (int i = 0; i < pageCount; i++) {

                BufferedImage image = pdf.renderImageWithDPI(i, 300);
                File imageFile = new File(originPath + File.separator + i + ".jpg");
                ImageIO.write(image, "PNG", imageFile);

                //调用图片识别方法
                String result = RapidOcrUtils.recognizeImage(imageFile, engine, paramConfig);

                if(imageFile.exists()){
                    imageFile.delete();
                }
                list.add(result);
            }

            document.close();

            if (tempFile.exists()){
                tempFile.delete();
            }


        }
        //照片
        else if(imageNames.contains(endFileName)){

            new File(originPath).mkdirs();
            File tempFile = new File(originPath + File.separator + file.getOriginalFilename());
            file.transferTo(tempFile);
            //调用识别方法
            String result = RapidOcrUtils.recognizeImage(tempFile, engine, paramConfig);

            if(tempFile.exists()){
                tempFile.delete();
            }
            list.add(result);

        }

        return list.toString().replaceAll("\\[|\\]", "");
    }

    @Override
    public IdCard recognizeID(MultipartFile frontFile , MultipartFile  backFile) throws IOException {

        String input = recognizeImage(frontFile);

        // 提取姓名
        String name = extractInfo(input, "姓名([^性]+)").trim();
        // 提取性别
        String gender = extractInfo(input, "性别([^民]+)").trim();
        // 提取民族
        String ethnicity = extractInfo(input, "民族([^出]+)").trim();
        // 提取出生日期
        String birthDate = extractInfo(input, "出生(\\d{4}年\\d{1,2}月\\d{1,2}日)([^住]+)").trim();
        // 提取住址
        String address = extractInfo(input, "住址([^公]+)");
        address = address.replaceAll(" ","");
        address = address.replaceAll("\\s+", "");
        // 提取身份证号码
        String idNumber = extractInfo(input, "[^0-9Xx]([1-9]\\d{16}[\\dXx])(?![0-9Xx])");

        //识别反面
        String backInput = recognizeImage(backFile);

        String issue = extractInfo(backInput, "签发机关([^有效]+)");
        String validDate = extractInfo(backInput, "有效期限\\s*(\\d{4}[./年-]?\\d{2}[./月-]?\\d{2}[日]?\\s*-\\s*\\d{4}[./年-]?\\d{2}[./月-]?\\d{2}[日]?|[\\d]{4}[./年-]?[\\d]{2}[./月-]?[\\d]{2}[日]?\\s*至\\s*[\\d]{4}[./年-]?[\\d]{2}[./月-]?[\\d]{2}[日]?|\\d{8}-\\d{4}\\.\\d{2}\\.\\d{2})");

        IdCard idCard = new IdCard();
        idCard.setName("null".equals(name)?null: name);
        idCard.setSex("null".equals(gender)?null: gender);
        idCard.setNation("null".equals(ethnicity)?null: ethnicity);
        idCard.setBirth("null".equals(birthDate)?null: birthDate);
        idCard.setAddress("null".equals(address)?null:  address);
        idCard.setNumber("null".equals(idNumber)?null: idNumber);
        idCard.setIssue("null".equals(issue)?null: issue);
        idCard.setValidDate("null".equals(validDate)?null: validDate);

        return idCard;

    }

    @Override
    public String recognizePDF(MultipartFile file) throws IOException {
        if(file != null){
            new File(originPath).mkdirs();
            File tempFile = new File(originPath + File.separator + file.getOriginalFilename());
            file.transferTo(tempFile);
            PDDocument document = PDDocument.load(tempFile);
            PDFTextStripper stripper = new PDFTextStripper();
            stripper.setSortByPosition(true);
            String text = stripper.getText(document);
            document.close();
            System.out.println(text);
            if (tempFile.exists()){
                tempFile.delete();
            }
            return text;
        }
        return "";
    }

    private static String extractInfo(String input, String regex) {
        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(input);
        if (matcher.find()) {
            return matcher.group(1);
        }
        return "null";
    }

}
